import numpy as np
from ultralytics import YOLO

classes = {0:'ripe',1:'half-ripe',2:'raw'}

#加载模型
model = YOLO('best.pt')
model(np.random.rand(100,100,3))

def predict_yolo(image):
    results = model(image)
    boxes = results[0].boxes
    # 按置信度排序，取最高的那个
    max_conf_idx = boxes.conf.argmax().item()  # 获取最大置信度的索引
    class_id = int(boxes.cls[max_conf_idx])  # 获取对应类别的ID
    class_name = model.names[class_id]  # 获取类别名称
    return class_name